from __future__ import annotations

import os
from typing import Optional, OrderedDict
import numpy as np

from common.geometry.bbox import nudge_bbox
from common.typing_utils import numeric
from common.logger import logger
from common.config import DEFAULT_CACHE_DIR
from common.io_utils.dirs import mkdir_p

from common.geometry.geom import Geom
from common.geometry.srs import coords_to_epsg_under_utm
from common.geometry.tools import project_geom
from common.geoimage.raster_dataset import RasterDataset
from common.geoimage.gdal_sys import warp_ds, gdal_env

from catalog.models.base_stac_item import BandDefinition

# NOTE: currently only support two remote TMS/WMS services
# and notice both of them can only be used with caution for scientific study only due to their non-commercial license.
# we could add more services in the future, especially services with more commercial friendly terms.
bands = OrderedDict()
bands["R"] = BandDefinition("Red", "Red", 0, 0.5, 0, 1, 0, np.uint8, None)
bands["G"] = BandDefinition("Green", "Green", 0, 0.5, 0, 1, 0, np.uint8, None)
bands["B"] = BandDefinition("Blue", "Blue", 0, 0.5, 0, 1, 0, np.uint8, None)
SUPPORTED_REMOTE_COLLECTIONS = {
    "GOOGLEMAPS/SATELLITE": {
        "href": "frmt_googlemaps_satellite",
        "default_resolutions": {"meter": 0.5, "degree": 4.5454545454545455e-06},
        "bands": bands,
        "supported_vis": ["RGB"],
    },
    "BINGMAPS/SATELLITE": {
        "href": "frmt_bingmaps_satellite",
        "default_resolutions": {"meter": 0.5, "degree": 4.5454545454545455e-06},
        "bands": bands,
        "supported_vis": ["RGB"],
    },
}


DEFAULT_WMS_CACHE_DIR = os.path.join(
    DEFAULT_CACHE_DIR, "catalog", "raster", "raster_remote_service"
)
mkdir_p(DEFAULT_WMS_CACHE_DIR)


def retrieve_by_query(
    collection_id: str,
    spatial: Geom,
    dest_srs: Optional[str] = None,
    dest_res: Optional[numeric] = None,
    fpath_dest: Optional[str] = None,
    resample_alg: str = "near",
    **kwargs,
) -> Optional[RasterDataset]:
    """Retrieve a raster dataset from remote service by a given spatial query.

    Parameters
    ----------
    collection_id : str
        Data collection id for the data of interest, e.g "GOOGLEMAPS/SATELLITE".
    spatial : Geom
        Spatial query as a `Geom` object that intersects with items of interest.
    dest_srs : Optional[str], optional
        Destination spatial reference system represented by an EPSG string, e.g. "EPSG:4326", "EPSG:4857", etc.
        by default None to use the item's own spatial reference system.
    dest_res : Optional[numeric], optional
        Destination spatial resolution, by default None to use the first item's own spatial resolution.
    fpath_dest : Optional[str], optional
        File path to the retrieved raster,
        by default None so not to save a raster file but return the raster as a `RasterDataset` object.
    resample_alg : str, optional
        Algorithm for resampling, if resampling is needed, by default "near".
        See <https://gdal.org/programs/gdalwarp.html#cmdoption-gdalwarp-r> for more algorithms supported.

    Returns
    -------
    Optional[RasterDataset]
        A `RasterDataset` object for the data retrieved,
        or `None` when a result raster file of geotiff or jpeg2000 format is requested to saved.

    Raises
    ------
    ValueError
        When an unsupported collection id is given.
    """
    if collection_id not in SUPPORTED_REMOTE_COLLECTIONS:
        raise ValueError(f"Unsupported collection_id: {collection_id}")

    src_spec = os.path.join(
        os.path.dirname(os.path.abspath(__file__)),
        "data",
        SUPPORTED_REMOTE_COLLECTIONS[collection_id]["href"] + ".xml",
    )

    if dest_srs is None:
        lon, lat = spatial.centroid._shapely_geom.x, spatial.centroid._shapely_geom.y
        dest_srs = "EPSG:{}".format(coords_to_epsg_under_utm(lon, lat))
        logger.info(
            f"No `dest_srs` is given, use '{dest_srs}' as the best UTM zone for the given `spatial`."
        )

    dest_extent = project_geom(spatial, "EPSG:4326", dest_srs).bounds

    if dest_res is None:
        dest_res = min(dest_extent.xsize / 4096, dest_extent.ysize / 4096)
    else:
        dest_extent = nudge_bbox(dest_extent, dest_res)

    with gdal_env(GDAL_DEFAULT_WMS_CACHE_PATH=DEFAULT_WMS_CACHE_DIR):
        if fpath_dest is not None:
            warp_ds(
                src_spec,
                dest_srs=dest_srs,
                dest_extent=dest_extent,
                dest_res=dest_res,
                resample_alg=resample_alg,
                fpath_dest=fpath_dest,
            )
        else:
            with warp_ds(
                src_spec,
                dest_srs=dest_srs,
                dest_extent=dest_extent,
                dest_res=dest_res,
                resample_alg=resample_alg,
            ) as gdal_raster:
                return RasterDataset.from_gdal_raster(gdal_raster)
